#include <rclcpp/rclcpp.hpp>
#include <stdlib.h>

#include <chrono>
#include <iostream>
#include <random>
#include <memory>

#include "behavior_planner/behavior_server_ros.h"
#include "common/basics/tic_toc.h"
#include "forward_simulator/multimodal_forward.h"
#include "forward_simulator/onlane_forward_simulation.h"
#include "semantic_map_manager/ros_adapter.h"
#include "semantic_map_manager/semantic_map_manager.h"
#include "semantic_map_manager/visualizer.h"
#include "sensor_msgs/msg/joy.hpp"
#include "vehicle_msgs/msg/control_signal.hpp"
// #include "vehicle_msgs/encoder.h"  // TODO: Update for ROS2

class OnlaneAiAgentNode : public rclcpp::Node {
public:
  OnlaneAiAgentNode() : Node("onlane_ai_agent") {
    this->declare_parameter<int>("ego_id", 0);
    this->declare_parameter<std::string>("agent_config_path", "");
    this->declare_parameter<double>("desired_vel", 6.0);
    this->declare_parameter<int>("autonomous_level", 2);
    this->declare_parameter<int>("aggressiveness_level", 3);
    
    ego_id = this->get_parameter("ego_id").as_int();
    agent_config_path = this->get_parameter("agent_config_path").as_string();
    desired_vel = this->get_parameter("desired_vel").as_double();
    int autonomous_level = this->get_parameter("autonomous_level").as_int();
    aggressiveness_level = this->get_parameter("aggressiveness_level").as_int();
    
    if (agent_config_path.empty()) {
      RCLCPP_ERROR(this->get_logger(), "Failed to get param agent_config_path");
      return;
    }
    
    ctrl_signal_pub_ = this->create_publisher<vehicle_msgs::msg::ControlSignal>("ctrl", 10);
    
    // Initialize components
    init_components(autonomous_level);
    
    // Create timer for control loop
    timer_ = this->create_wall_timer(
      std::chrono::duration<double>(1.0 / fs_work_rate),
      std::bind(&OnlaneAiAgentNode::control_callback, this)
    );
    
    // Create timer for visualization
    vis_timer_ = this->create_wall_timer(
      std::chrono::duration<double>(1.0 / visualization_msg_rate),
      std::bind(&OnlaneAiAgentNode::visualization_callback, this)
    );
  }
  
  ~OnlaneAiAgentNode() {
    if (p_bp_server_) delete p_bp_server_;
    if (p_smm_vis_) delete p_smm_vis_;
    if (p_ctrl_input_smm_buff_) delete p_ctrl_input_smm_buff_;
  }

private:
  // ROS2 components
  rclcpp::Publisher<vehicle_msgs::msg::ControlSignal>::SharedPtr ctrl_signal_pub_;
  rclcpp::TimerBase::SharedPtr timer_;
  rclcpp::TimerBase::SharedPtr vis_timer_;
  
  // Parameters
  int ego_id;
  std::string agent_config_path;
  double desired_vel;
  int aggressiveness_level = 3;
  
  // Core components (keeping existing logic)
  planning::BehaviorPlannerServer* p_bp_server_{nullptr};
  moodycamel::ReaderWriterQueue<semantic_map_manager::SemanticMapManager>*
      p_ctrl_input_smm_buff_{nullptr};
  semantic_map_manager::Visualizer* p_smm_vis_{nullptr};
  
  common::State desired_state;
  bool has_init_state = false;
  semantic_map_manager::SemanticMapManager last_smm;
  planning::OnLaneForwardSimulation::Param sim_param;
  
  // Random behavior
  std::mt19937 rng;
  double vel_noise = 0.0;
  int cnt = 0;
  
  void init_components(int autonomous_level);

  int semantic_map_update_callback(
      const semantic_map_manager::SemanticMapManager& smm) {
    if (p_bp_server_) p_bp_server_->PushSemanticMap(smm);
    
    if (!has_init_state) {
      desired_state = smm.ego_vehicle().state();
      has_init_state = true;
      desired_state.print();
    }
    return 0;
  }
  
  int behavior_update_callback(
      const semantic_map_manager::SemanticMapManager& smm) {
    if (p_ctrl_input_smm_buff_) p_ctrl_input_smm_buff_->try_enqueue(smm);
    return 0;
  }

  void random_behavior() {
    if (cnt == 0) {
      std::uniform_real_distribution<double> dist_vel(-2, 5);
      vel_noise = dist_vel(rng);
      p_bp_server_->set_user_desired_velocity(desired_vel + vel_noise);
      RCLCPP_INFO(this->get_logger(), "[OnlaneAi]%d - desired velocity: %lf", 
                  ego_id, desired_vel + vel_noise);
    }
    cnt++;
    if (cnt >= 2000) cnt = 0;
  }

  void control_callback() {
    if (!has_init_state) return;
    if (p_bp_server_ == nullptr) return;
    if (p_ctrl_input_smm_buff_ == nullptr) return;
    
    bool is_map_updated = false;
    decimal_t previous_stamp = last_smm.time_stamp();
    while (p_ctrl_input_smm_buff_->try_dequeue(last_smm)) {
      is_map_updated = true;
    }
    if (!is_map_updated) return;
    
    decimal_t delta_t = last_smm.time_stamp() - previous_stamp;
    if (delta_t > 100.0 / fs_work_rate) delta_t = 1.0 / fs_work_rate;
    
    decimal_t command_vel = p_bp_server_->reference_desired_velocity();
    decimal_t speed_limit;
    if (last_smm.GetSpeedLimit(last_smm.ego_vehicle().state(),
                               last_smm.ego_behavior().ref_lane,
                               &speed_limit) == kSuccess) {
      command_vel = std::min(speed_limit, command_vel);
    }
    common::Vehicle ego_vehicle = last_smm.ego_vehicle();
    ego_vehicle.set_state(desired_state);
    sim_param.idm_param.kDesiredVelocity = command_vel;
    
    common::Vehicle leading_vehicle;
    common::State state;
    decimal_t distance_residual_ratio = 0.0;
    const decimal_t lat_range = 2.2;
    last_smm.GetLeadingVehicleOnLane(last_smm.ego_behavior().ref_lane,
                                     desired_state,
                                     last_smm.surrounding_vehicles(), lat_range,
                                     &leading_vehicle, &distance_residual_ratio);
    if (planning::OnLaneForwardSimulation::PropagateOnce(
            common::StateTransformer(last_smm.ego_behavior().ref_lane),
            ego_vehicle, leading_vehicle, delta_t, sim_param,
            &state) != kSuccess) {
      RCLCPP_ERROR(this->get_logger(), "[AiAgent]Err-Simulation error (with leading vehicle).");
      return;
    }
    
    common::VehicleControlSignal ctrl(state);
    {
      vehicle_msgs::msg::ControlSignal ctrl_msg;
      // TODO: Update encoder for ROS2
      // vehicle_msgs::Encoder::GetRosControlSignalFromControlSignal(
      //     ctrl, this->now(), std::string("map"), &ctrl_msg);
      ctrl_signal_pub_->publish(ctrl_msg);
    }
    desired_state = ctrl.state;
  }

  void visualization_callback() {
    if (!has_init_state) return;
    if (p_smm_vis_ == nullptr) return;
    
    rclcpp::Time tnow = this->now();
    p_smm_vis_->VisualizeDataWithStamp(tnow, last_smm);
    p_smm_vis_->SendTfWithStamp(tnow, last_smm);
  }
  
  // Constants
  static constexpr double fs_work_rate = 50.0;
  static constexpr double visualization_msg_rate = 20.0;
  static constexpr double bp_work_rate = 20.0;
};

void OnlaneAiAgentNode::init_components(int autonomous_level) {
  // Get desired velocity noise
  rng.seed(
      std::chrono::high_resolution_clock::now().time_since_epoch().count());
  planning::MultiModalForward::ParamLookUp(aggressiveness_level, &sim_param);
  RCLCPP_INFO(this->get_logger(), "[OnlaneAi]%d - aggresive: %d", ego_id, aggressiveness_level);
  
  // Declare smm (TODO: Update for ROS2 when semantic_map_manager is migrated)
  // semantic_map_manager::SemanticMapManager semantic_map_manager(
  //     ego_id, agent_config_path);
  // semantic_map_manager::RosAdapter smm_ros_adapter(shared_from_this(), &semantic_map_manager);
  // p_smm_vis_ = new semantic_map_manager::Visualizer(shared_from_this(), ego_id);
  
  // Declare bp (TODO: Update for ROS2 when behavior_planner is migrated)
  // p_bp_server_ = new planning::BehaviorPlannerServer(shared_from_this(), bp_work_rate, ego_id);
  // p_bp_server_->set_user_desired_velocity(desired_vel);
  // p_bp_server_->set_autonomous_level(autonomous_level);
  // p_bp_server_->set_aggressive_level(aggressiveness_level);
  // p_bp_server_->enable_hmi_interface();
  
  // smm_ros_adapter.BindMapUpdateCallback(
  //     std::bind(&OnlaneAiAgentNode::semantic_map_update_callback, this, std::placeholders::_1));
  // p_bp_server_->BindBehaviorUpdateCallback(
  //     std::bind(&OnlaneAiAgentNode::behavior_update_callback, this, std::placeholders::_1));
  
  // smm_ros_adapter.Init();
  // p_bp_server_->Init();
  
  p_ctrl_input_smm_buff_ = new moodycamel::ReaderWriterQueue<
      semantic_map_manager::SemanticMapManager>(100);
  
  // p_bp_server_->Start();
}

int main(int argc, char** argv) {
  rclcpp::init(argc, argv);
  auto node = std::make_shared<OnlaneAiAgentNode>();
  
  // TODO: Start bp_server when behavior_planner is migrated to ROS2
  // node->start_behavior_planner();
  
  rclcpp::spin(node);
  rclcpp::shutdown();
  return 0;
}
